'''
Created on Jan 25, 2010

@author: mkiyer
'''

import logging
import os
import sys
import subprocess
import shutil

from veggie.genome.chrom import HG18


def filter_references(insamfile, outsamfile,
                      references=None):    
    if (references == None):
        return
    logger = logging.getLogger(__name__)
    outsam = open(outsamfile, "w")
    insam = open(insamfile, "r")
    for header_line in insam:
        # once we find the first line that does not start with '@' the header is over
        if not header_line.startswith("@"):
            break
        header_fields = header_line.strip().split('\t')
        if header_fields[0] == '@SQ':
            fields = dict([field.split(':', 1) for field in header_fields[1:]])
            assert 'SN' in fields
            if fields['SN'] not in references:
                logger.debug('filter_references: rejected reference %s' % fields['SN'])
                continue        
        outsam.write(header_line)
    # still have to deal with the first read before we can start a new loop
    rejected_reads = 0    
    getrname = lambda r: r.split('\t', 3)[2]
    read_line = header_line
    rname = getrname(read_line)
    if rname == '*':
        outsam.write(read_line)
    elif rname in references:
        outsam.write(read_line)
    else:
        rejected_reads += 1
    # loop through rest of reads
    for read_line in insam:
        rname = getrname(read_line)
        if rname == '*':
            outsam.write(read_line)
        elif rname in references:
            outsam.write(read_line)
        else:
            rejected_reads += 1    
    insam.close()
    outsam.close()
    logger.debug('filter_references: removed %d reads from non-matching references' % (rejected_reads))
    return outsamfile

def convert_sam_to_bam(insamfiles, bamfile, tmp_path=None, references=None):
    if references == None:
        #TODO: customize to any version of genome
        references = HG18.chrom_sizes
    logger = logging.getLogger(__name__)
    tmpfiles = []
    mergebamfiles = []
    if tmp_path == None:
        tmp_path = os.path.abspath(os.path.dirname(bamfile))
    # first have to convert the SAM files to BAM and prepare them
    for insamfile in insamfiles:
        # remove references that are contamination
        myfiltersam = os.path.join(tmp_path, os.path.splitext(os.path.basename(insamfile))[0] + '.filter.sam')
        tmpfiles.append(myfiltersam)
        logger.debug("Filtering references from SAM %s -> SAM %s" % (insamfile, myfiltersam))
        filter_references(insamfile, myfiltersam, references=references)
        nextinfile = myfiltersam
        # convert to BAM
        mybamfile = os.path.splitext(nextinfile)[0] + '.bam'
        tmpfiles.append(mybamfile)
        logger.debug("Converting SAM %s -> BAM %s" % (nextinfile, mybamfile))
        subprocess.call(['samtools', 'view', '-S', nextinfile, '-b', '-o', mybamfile])
        nextinfile = mybamfile 
        # sort
        mysortedbam = os.path.splitext(nextinfile)[0] + '.sorted.bam'
        tmpfiles.append(mysortedbam)
        logger.debug("Sorting BAM %s -> %s" % (nextinfile, mysortedbam))
        # sort automatically adds ".bam" to the file, so remove the extension
        subprocess.call(['samtools', 'sort', nextinfile, os.path.splitext(mysortedbam)[0]])
        # index        
        logger.debug("Indexing Sorting BAM %s" % (mysortedbam))
        subprocess.call(['samtools', 'index', mysortedbam])
        tmpfiles.append(mysortedbam + '.bai')
        nextinfile = mysortedbam
        # remove duplicates
        myrmdupbam = os.path.splitext(nextinfile)[0] + '.rmdup.bam'
        tmpfiles.append(myrmdupbam)
        logger.debug("Removing duplicates from BAM %s -> %s" % (nextinfile, myrmdupbam))
        subprocess.call(['samtools', 'rmdup', '-s', nextinfile, myrmdupbam])
        nextinfile = myrmdupbam
        # files to merge
        mergebamfiles.append(nextinfile)        
    # merge
    if len(mergebamfiles) == 1:
        # just one file, so copy it rather than merge it
        # it is already sorted too.. so just re-index it
        logger.debug("Copying BAM file %s -> %s" % (mergebamfiles[0], bamfile))
        shutil.copyfile(mergebamfiles[0], bamfile)
    else:
        #mymergedbam = os.path.splitext(bamfile)[0] + '.merge.bam'
        #tmpfiles.append(mymergedbam)
        merge_cmd = ['samtools', 'merge', bamfile]
        merge_cmd.extend(mergebamfiles)
        logger.debug("Merging BAM files %s -> %s" % (mergebamfiles, bamfile))
        subprocess.call(merge_cmd)
        # sort
        #logger.debug("Sorting BAM file %s -> %s" % (mymergedbam, bamfile))
        # sort automatically adds ".bam" to the file, so remove the extension
        #subprocess.call(['samtools', 'sort', mymergedbam, os.path.splitext(bamfile)[0]])
    logger.debug("Indexing BAM file %s" % (bamfile))
    subprocess.call(['samtools', 'index', bamfile])
    # remove temporary files
    for tmpfile in tmpfiles:
        logger.debug("Removing %s" % tmpfile)
        os.remove(tmpfile)
        
if __name__ == '__main__':
    convert_sam_to_bam([sys.argv[1]], sys.argv[2])
